from io import StringIO
import sys
from unittest.mock import patch, MagicMock
from malwarebazaar import malwarebazaar
import unittest


class TestMalwarebazaarMethods(unittest.TestCase):
    def test_main_missing_input(self):
        with patch('sys.stdout', new=StringIO()) as mock_cmd:
            sys.argv = ["cmd"]
            malwarebazaar.main()
            self.assertEqual(mock_cmd.getvalue(),
                             'ERROR: Input is not in proper JSON format\n')

    def test_main_success(self):
        with patch('sys.stdout', new=StringIO()) as mock_cmd:
            with patch('malwarebazaar.malwarebazaar.analyze',
                       new=MagicMock(return_value={'test': 'val'})) as mock:
                sys.argv = ["cmd", "input"]
                malwarebazaar.main()
                expected = '{"test": "val"}\n'
                self.assertEqual(mock_cmd.getvalue(), expected)
                mock.assert_called_once()

    def test_isInJson_tail_greater_than_max_depth(self):
        max_depth = 1000
        tail = 2000
        test_string = "helo"
        input_json = {
            "value": "test",
            "test": "value",
            "arr": ["Foo", "Bar", "Hello"],
            "dict1": {"key1": "val", "key2": "helo"}
        }
        self.assertEqual(malwarebazaar.isInJson(input_json, test_string, max_depth, tail), False)

    def test_isInJson_string_found_in_dict(self):
        test_string = "helo"
        input_json = {
            "value": "test",
            "test": "value",
            "arr": ["Foo", "Bar", "Hello"],
            "dict1": {"key1": "val", "key2": "helo"}
        }
        self.assertEqual(malwarebazaar.isInJson(input_json, test_string), True)

    def test_isInJson_dict_in_list(self):
        max_depth = 1000
        tail = 1
        test_string = "helo"
        input_json = {
            "key1": "test",
            "key2": "value",
            "key3": ["Foo", "Bar", "Hello"],
            "nested_list": [{"key1": "val", "key2": "helo"}]
        }
        self.assertEqual(malwarebazaar.isInJson(input_json, test_string, max_depth, tail), True)

    def test_isInJson_string_found_in_arr(self):
        test_string = "helo"
        input_json = {
            "value": "test",
            "test": "value",
            "arr": ["Foo", "Bar", "helo"],
            "dict1": {"Hello": "val", "key": "val"}
        }
        self.assertEqual(malwarebazaar.isInJson(input_json, test_string), True)

    def test_isInJson_string_not_found(self):
        test_string = "ValNotInJSON"
        input_json = {
            "value": "test",
            "test": "value",
            "arr": ["Foo", "Bar", "helo"],
            "dict1": {"Hello": "val", "key": "val"}
        }
        self.assertEqual(malwarebazaar.isInJson(input_json, test_string), False)

    def test_analyze(self):
        """simulated sendReq and prepareResults with 2 mock objects
            and variables sendReqOutput and prep_res_sim,
            input created for analyze method call
            and then we compared results['summary'] with 'no result' """
        sendReqOutput = {'threat': 'no_result', "query_status": "ok",
                         'data': [{'sha256_hash': 'notavalidhash'}]}
        input = '{"artifactType": "hash", "value": "1234"}'
        input2 = '{"artifactType": "tlsh", "value": "1234"}'
        input3 = '{"artifactType": "gimphash", "value": "1234"}'
        prep_res_sim = {'response': '',
                        'summary': 'no result', 'status': 'info'}

        with patch('malwarebazaar.malwarebazaar.sendReq',
                   new=MagicMock(return_value=sendReqOutput)) as mock:
            with patch('malwarebazaar.malwarebazaar.prepareResults',
                       new=MagicMock(return_value=prep_res_sim)) as mock2:
                results = malwarebazaar.analyze(input)
                results2 = malwarebazaar.analyze(input2)
                results3 = malwarebazaar.analyze(input3)
                self.assertEqual(results["summary"], prep_res_sim['summary'])
                self.assertEqual(results2["summary"], prep_res_sim['summary'])
                self.assertEqual(results3["summary"], prep_res_sim['summary'])
                self.assertEqual(results["status"], "info")
                self.assertEqual(results2["status"], "info")
                self.assertEqual(results3["status"], "info")
                mock2.assert_called()
                mock.assert_called()

    def test_analyze_result(self):
        """simulated sendReq and prepareResults with 2 mock objects
            and variables sendReqOutput and prep_res_sim,
            input created for analyze method call
            and then we compared results['summary'] with 'no result' """
        sendReqOutput = {'threat': 'threat', "query_status": "notok", 'data': [
            {'sha256_hash': 'validhash'}]}
        input = '{"artifactType": "hash", "value": "1234"}'
        input2 = '{"artifactType": "tlsh", "value": "1234"}'
        input3 = '{"artifactType": "gimphash", "value": "1234"}'
        prep_res_sim = {'response': '',
                        'summary': 'Bad', 'status': 'threat'}

        with patch('malwarebazaar.malwarebazaar.sendReq',
                   new=MagicMock(return_value=sendReqOutput)) as mock:
            with patch('malwarebazaar.malwarebazaar.prepareResults',
                       new=MagicMock(return_value=prep_res_sim)) as mock2:
                results = malwarebazaar.analyze(input)
                results2 = malwarebazaar.analyze(input2)
                results3 = malwarebazaar.analyze(input3)
                self.assertEqual(results["summary"], prep_res_sim['summary'])
                self.assertEqual(results2["summary"], prep_res_sim['summary'])
                self.assertEqual(results3["summary"], prep_res_sim['summary'])
                self.assertEqual(results["status"], "threat")
                self.assertEqual(results2["status"], "threat")
                self.assertEqual(results3["status"], "threat")
                mock2.assert_called()
                mock.assert_called()

    def test_prepareResults_illegal_search_term(self):
        # illegal search term
        raw = {'query_status': 'illegal_search_term'}
        expected = {'response': raw, 'status': 'info', 'summary': 'no result'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_prepareResults_empty(self):
        # raw is empty
        raw = {}
        expected = {'response': raw, 'status': 'caution',
                    'summary': 'internal_failure'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_prepareResults_threat(self):
        raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash',
                                               'vendor_intel':
                                                   {'ReversingLabs':
                                                       {'status':
                                                           'MALICIOUS'}},
                                               'signature': 'abcd1234',
                                               'tags': ['tag1']}]}
        expected = {'response': raw, 'status': 'threat', 'summary': 'abcd1234'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_prepareResults_caution(self):
        # raw is empty
        raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash',
                                               'vendor_intel':
                                                   {'Triage': {'score': '6'}},
                                                   'signature': 'abcd1234',
                                                   'tags': ['tag1']}]}
        expected = {'response': raw,
                    'status': 'caution', 'summary': 'abcd1234'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_prepareResults_info(self):
        # raw is empty
        raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash',
                                               'vendor_intel':
                                                   {'Triage': {'score': '3'}},
                                                   'signature': 'abcd1234',
                                                   'tags': ['tag1']}]}
        expected = {'response': raw, 'status': 'info', 'summary': 'abcd1234'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_prepareResults_ok(self):
        # raw is empty
        raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash',
                                               'vendor_intel':
                                                   {'Triage': {'score': '1'}},
                                                   'signature': 'abcd1234',
                                                   'tags': ['tag1']}]}
        expected = {'response': raw, 'status': 'ok', 'summary': 'abcd1234'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_prepareResults_ok_tags(self):
        # raw is empty
        raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash',
                                               'vendor_intel':
                                                   {'Triage': {'score': '1'}},
                                                   'tags': ['tag1']}]}
        expected = {'response': raw, 'status': 'ok', 'summary': 'tag1'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_prepareResults_ok_yomi(self):
        # raw is empty
        raw = {'query_status': 'ok',
               'data': [{'sha256_hash': 'validhash',
                        'vendor_intel':
                            {'YOROI_YOMI':
                                {'detection':
                                    'detection1',
                                    'summary': '0.1'}}}]}
        expected = {'response': raw, 'status': 'ok', 'summary': 'detection1'}
        results = malwarebazaar.prepareResults(raw)
        self.assertEqual(results, expected)

    def test_buildReqGimqhash(self):
        result = malwarebazaar.buildReq('gimphash', '')
        self.assertEqual(
            result, {'query': 'get_gimphash', 'gimphash': ''})

    def test_buildReqHash(self):
        result = malwarebazaar.buildReq('hash', '')
        self.assertEqual(
            result, {'query': 'get_info', 'hash': ''})

    def test_buildReqtlshhash(self):
        result = malwarebazaar.buildReq('tlsh', '')
        self.assertEqual(
            result, {'query': 'get_tlsh', 'tlsh': ''})

    # simulate API response and makes sure sendReq gives a response,
    # we are just checking if sendReq gives back anything
    def test_sendReq(self):
        with patch('requests.post',
                   new=MagicMock(return_value=MagicMock())) as mock:
            response = malwarebazaar.sendReq(
                {'baseUrl': 'https://www.randurl.xyz'}, 'example_data')
            self.assertIsNotNone(response)
            mock.assert_called_once()
